import os
from nlp_tools.corpus.classify.dialogue_corpus_loader import ClassifyNerCorpusLoader
from nlp_tools.metrics.classification import F1CategoryCallback
from nlp_tools.tasks.classification import ClassificationCnn
from nlp_tools.tasks.classification.cnn_attention_model import CNN_Attention_Model
from nlp_tools.tasks.classification.entity_level_model import ClassificationEntityLevel
from nlp_tools.processors.dialogue_sequence_processor_huggingface import DialogueSequenceProcessor
from nlp_tools.processors.classification.classification_label_processor import ClassificationLabelProcessor
from nlp_tools.tokenizer.hugging_tokenizer import HuggingTokenizer
from nlp_tools.embeddings.hugginface.autoembedding import AutoEmbedding

from nlp_tools.callbacks.classification.f1score_save_callback import F1SaveCallback

import random
import numpy as np
import tensorflow as tf

def seed_tensorflow(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)
    os.environ['TF_DETERMINISTIC_OPS'] = '1' # pip install tensorflow-determinism

seed_tensorflow(2021)

model_save_path = '/home/qiufengfeng/nlp/train_models/bad_cargonName/'

train_data = ClassifyNerCorpusLoader().load_data("/home/fanfanfeng/working/contraband/train.csv")
valid_data = ClassifyNerCorpusLoader().load_data("/home/fanfanfeng/working/contraband/test.csv")

all_data = train_data+ valid_data
from sklearn.model_selection import train_test_split
train_data,valid_data = train_test_split(all_data,test_size=0.15)



bert_model_path = "hfl/chinese-bert-wwm-ext"
label_list = ["不是","是"]
label_dict = {key:index for index,key in enumerate(label_list)}



text_tokenizer  = HuggingTokenizer(bert_model_path)

from nlp_tools.embeddings.hugginface.prefix_prompt.auto_embedding_prefix_embedding import AutoEmbeddingPrefixEmbedding
embedding  = AutoEmbedding(bert_model_path,text_tokenizer.tokenizer.model_input_names)

# 默认是不需要分词或者对训练数据进行处理的，如果需要，则要重写text_tokenizer和相应的processor
sequenceProcessor = DialogueSequenceProcessor(text_tokenizer=text_tokenizer,return_ner_masking=True)
labelProcessor = ClassificationLabelProcessor(vocab2idx=label_dict,multi_label=False)
model = ClassificationEntityLevel(
    embedding=embedding,
    text_processor=sequenceProcessor,
    label_processor=labelProcessor,
    use_rdrop=True,
    use_FGM=False,
    max_sequence_length=128)

model_save_path = 'contraband/tf_models/train_model_precision_random/'
precision_callback = F1SaveCallback(model,model_save_path,valid_data,label_names=label_list,score_type="precision_score")

model_save_path1 = 'contraband/tf_models/train_model_f1_random/'
f1_callback = F1SaveCallback(model,model_save_path1,valid_data,label_names=label_list)

from keras.api._v2.keras.callbacks import  TensorBoard


tensorboad_callback = TensorBoard(log_dir=os.path.join(model_save_path,'tensorboard_logs'),write_graph=True,write_images=True)
model.fit(train_data,validate_data=valid_data,epochs=60,callbacks=[f1_callback,tensorboad_callback,precision_callback],batch_size=25)

